# -*- coding: utf-8 -*-
import mistune

from pygments import highlight
from pygments.lexers import get_lexer_by_name
from pygments.formatters import html

class CustomRenderer(mistune.Renderer):

    def block_code(self, code, language=None):
        if not language:
            return "\n<pre><code>{}</code></pre>\n".format(mistune.escape(code))
        lexer = get_lexer_by_name(language, stripall=True)
        formatter = html.HtmlFormatter()
        return highlight(code, lexer, formatter)

def unittest():
    myRenderer = CustomRenderer()
    markdown = mistune.Markdown(renderer=myRenderer)
    
    with open("test.html", 'w') as fd:
        html = markdown("```python\nassert 1 == 1\n```")
        fd.write(html)

if __name__ == '__main__':
    unittest()
